module Hexena.Cmds.EncapsulateEXE (run) where

import Data.Array.IArray (elems)
import qualified Codec.Compression.BZip as CBZ
import qualified Codec.Compression.GZip as CGZ
import Data.Binary.Get (runGet)
import Data.Binary.Put
import Data.Char (ord)
import qualified Data.Digest.MD5 as MD5
import qualified Data.Digest.SHA1 as SHA1
import Data.IFF
import Data.Maybe (mapMaybe)
import System.Console.GetOpt
import System.Environment (getArgs)
import System.IO (hPutStrLn, stderr)
import System.Posix.Files (getFileStatus, fileID, fileSize)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL

import Hexena.Chunks
import Hexena.MZHeader
import Hexena.PEHeader
import Hexena.Math.Statistics

data CLIOptions = CLIOptions { optCompression :: Int
                             , optJUNK :: Bool
                             , optKeepName :: Bool
                             , optKeepBinary :: Bool
                             , optAnnotation :: String
                             }
                | CLIOptHelp
                  deriving Show

chooseCompression :: Int -> B.ByteString -> B.ByteString
chooseCompression 0 = id
chooseCompression 1 = rleCompress
chooseCompression 2 = \x -> B.concat $ BL.toChunks $ CGZ.compress $ BL.fromChunks [x]
chooseCompression 3 = \x -> B.concat $ BL.toChunks $ CBZ.compress $ BL.fromChunks [x]
chooseCompression _ = error "Illegal compression level"

rleCompress :: B.ByteString -> B.ByteString
rleCompress = B.concat . map (B.concat . BL.toChunks . mkBS) . doSpan . B.group
    where doSpan = doRLE . span (\x -> B.length x == 1)
          doRLE ([], rs) = handleRun rs
          doRLE (ss, rs) = (length ss, B.concat ss) : handleRun rs
          handleRun [] = []
          handleRun (x:xs) = (negate $ B.length x, x) : doSpan xs
          mkBS :: (Int, B.ByteString) -> BL.ByteString
          mkBS (l,b) | l < 0 = runPut $ putRun l b
                     | otherwise = runPut $ putDat l b
          putRun l b | l >= -128 = do putWord8 $ fromIntegral l
                                      putWord8 $ B.head b
                     | otherwise = do putWord8 (-128)
                                      putWord8 $ B.head b
                                      putRun (l + 128) b
          putDat l b | l < 127 = do putWord8 $ fromIntegral l
                                    putByteString b
                     | otherwise = do putWord8 127
                                      putByteString $ B.take 127 b
                                      putDat (l - 127) $ B.drop 127 b

options :: [ OptDescr (CLIOptions -> CLIOptions) ]
options = [ Option "h?" ["help"] (NoArg (const CLIOptHelp)) "Show Help Message"
          , Option "c" ["compression"] (ReqArg (\arg optval -> optval { optCompression = read arg } ) "[0..3]") "compression level"
          , Option "a" ["annotation"] (ReqArg (\arg optval -> optval { optAnnotation = arg } ) "ANNO") "Add some file annotation"
          , Option "J" ["JUNK"] (NoArg (\o -> o { optJUNK = True })) "Add JUNK Chunk"
          , Option "" ["keep-filename"] (NoArg (\o -> o { optKeepName = True })) "Keep original filename in malware encapsulation file"
          , Option "k" ["kill-binary"] (NoArg (\o -> o { optKeepBinary = False })) "Kill original binary in malware encapsulation file -- no binary will be included"
          ]

header :: String
header = "encapsulateEXE"

parseCLI :: IO (CLIOptions, [String])
parseCLI = let oDefault = CLIOptions 2 False False True ""
               foldFun :: (CLIOptions -> CLIOptions) -> CLIOptions -> CLIOptions
               foldFun _ CLIOptHelp = CLIOptHelp
               foldFun x y = x y
               pFun (opts, files, []) = return (foldr foldFun oDefault opts, files)
               pFun (_, _, errs) = ioError (userError (concat errs ++ usageInfo header options))
           in do argv <- getArgs
                 pFun $ getOpt Permute options argv


makeBODY :: CLIOptions -> B.ByteString -> Maybe T
makeBODY o | optKeepBinary o = Just . Cons (chunkIdFromString "BODY") . Chunk . compF
           | otherwise = const Nothing
    where compF = chooseCompression $ optCompression o

-- | TODO: specify format!
makeMWHD :: CLIOptions -> B.ByteString -> T
makeMWHD o d = Cons (chunkIdFromString "MWHD") $ Chunk $ B.concat $ BL.toChunks $ runPut $ do
                 putWord32be $ fromIntegral $ B.length d
                 putWord8 $ if optKeepBinary o then fromIntegral $ optCompression o else 0
                 putWord8 $ if optKeepBinary o then 1 else 0

makeJUNK :: Bool -> Maybe T
makeJUNK False = Nothing
makeJUNK True = Just $ Cons (chunkIdFromString "JUNK") (Chunk B.empty)

makeFNAM :: CLIOptions -> String -> Maybe T
makeFNAM opts fn | optKeepName opts = Just $ Cons (chunkIdFromString "FNAM") (Chunk $ B.pack $ map (fromIntegral . ord) fn)
                 | otherwise = Nothing

makeANNO :: CLIOptions -> Maybe T
makeANNO = f . optAnnotation
    where f [] = Nothing
          f x = Just $ saveANNO x

main2 :: CLIOptions -> String -> IO ()
main2 _ [] = putStrLn "Please provide file name to operate on!"
main2 CLIOptHelp _ = putStrLn $ usageInfo header options
main2 opt fname = do
  dat <- B.readFile fname
  stats <- getFileStatus fname
  let dat' = B.unpack dat
  hists <- calcBothHist dat'
  let md5 = MD5.hash dat'
  let sha1 = SHA1.hash dat'
  let (mzhdr, lfanew) = runGet checkFile $ BL.fromChunks [dat]
  let pehdr = readPEFrom $ BL.fromChunks [dat]
  let chunks = [ Just $ makeMWHD opt dat
               , makeANNO opt
               , makeBODY opt dat
               , makeFNAM opt fname
               , Just $ savePEHD pehdr
               , Just $ Cons (chunkIdFromString "MZRW") $ Chunk $ B.take 64 dat
               , Just $ makeMZHD mzhdr lfanew
               , makeJUNK (optJUNK opt)
               , Just $ uncurry makeFORMHIST hists
               , Just $ saveHASH md5 sha1
               ]
  let chunks' = mapMaybe id chunks
  putStrLn $ "\tsha1: " ++ show sha1
  putStrLn $ "\tmd5: " ++ show md5
  putStrLn $ "\t#histogram checks: " ++ show (sum $ elems $ fst hists) ++ ", " ++ show (sum $ elems $ snd hists)
  if B.take 2 dat == B.pack [0x4d, 0x5a] then
      writeIFF (fname ++ ".hexena") chunks'
      else
      hPutStrLn stderr $ "File '" ++ fname ++ "' does not seem to be a MZ file."
  putStrLn $ "\tChunks written: " ++ show (length chunks')

main1 :: CLIOptions -> [String] -> IO ()
main1 CLIOptHelp _ = putStrLn $ usageInfo header options
main1 _ [] = error "Error! Provide filenames!"
main1 o fs = mapM_ f fs
  where f fname = do putStrLn $ "Processing: " ++ fname
                     stats <- getFileStatus fname
                     putStrLn $ "\tFilesize: " ++ show (fileSize stats)
                     putStrLn $ "\tInode: " ++ show (fileID stats)
                     if fileSize stats < 6000001 then main2 o fname else putStrLn "\tWarning! Ignoring file due to size."

main :: IO ()
main = parseCLI >>= uncurry main1

run :: [String] -> IO ()
run fn = let oDefault = CLIOptions 2 False False True ""
         in main1 oDefault fn
